package com.chat.simplechat.interceptor;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.chat.simplechat.constant.WebSocketConstant;
import com.chat.simplechat.entity.LoginUser;
import com.chat.simplechat.jwt.TokenBuilder;
import com.chat.simplechat.service.UserInfoService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Map;
import java.util.Optional;

/**
 * WebSocket握手拦截器
 * @author wangyabin
 * @date 2022/10/28 17:37
 */
@Slf4j
@Component
public class WebSocketInterceptor implements HandshakeInterceptor {

    @Autowired
    private UserInfoService userInfoService; 
    
    /**
     * websocket无法设置header，只能通过子协议携带Sec-WebSocket-Protocol参数
     * 可以将token放入该字段中传入
     * @param request
     * @param response
     * @param wsHandler
     * @param attributes
     * @return
     * @throws Exception
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        if(request instanceof ServletServerHttpRequest){
            HttpServletRequest req = ((ServletServerHttpRequest) request).getServletRequest();
            String token = req.getHeader("Sec-WebSocket-Protocol");
            if(!StringUtils.hasText(token)){
                response.setStatusCode(HttpStatus.FORBIDDEN);
                log.info("[WebSocket]Token为空，身份验证失败!");
                return false;
            }
            Map<String,String> tokenMap = TokenBuilder.verifyToken(token);
            LoginUser user = JSONObject.parseObject(JSON.toJSONString(tokenMap),LoginUser.class);
            if(!Optional.ofNullable(user).isPresent()){
                response.setStatusCode(HttpStatus.FORBIDDEN);
                log.info("[WebSocket]用户不存在，身份验证失败!");
                return false;
            }
            log.info("[WebScoket]用户信息:{}",user);
            user = userInfoService.getLoginUser(user.getUsername());
            String userId = user.getUserId();
            attributes.put(WebSocketConstant.LOGIN_USER, user);
            attributes.put(WebSocketConstant.USER_ID,userId);
            attributes.put(HttpSessionHandshakeInterceptor.HTTP_SESSION_ID_ATTR_NAME, req.getSession().getId());
        }
        return true;
    }

    /**
     * 
     * @param request
     * @param response
     * @param wsHandler
     * @param exception
     */
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
        log.info("--------------------------websocket after handshake start--------------------------");
        HttpServletRequest httpRequest = ((ServletServerHttpRequest) request).getServletRequest();
        HttpServletResponse httpResponse = ((ServletServerHttpResponse) response).getServletResponse();
        if (StringUtils.hasText(httpRequest.getHeader("Sec-WebSocket-Protocol"))) {
            httpResponse.addHeader("Sec-WebSocket-Protocol", httpRequest.getHeader("Sec-WebSocket-Protocol"));
        }
        log.info("--------------------------websocket after handshake end--------------------------");
    }
}
